#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：t1.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/8 16:12
遗传算法求解FJSP
'''

import math
import random
import json
from dataclasses import dataclass
from typing import Dict, List, Tuple, Optional
import matplotlib.pyplot as plt

random.seed(42)


@dataclass
class Operation:
    # mapping: machine_id -> processing_time
    machine_time: Dict[int, int]


@dataclass
class Job:
    # operations in order
    ops: List[Operation]


@dataclass
class Problem:
    jobs: List[Job]
    m: int  # number of machines

    @property
    def n_jobs(self) -> int:
        return len(self.jobs)

    @property
    def n_ops_total(self) -> int:
        return sum(len(j.ops) for j in self.jobs)

    def op_choices(self, job_id: int, op_idx: int) -> List[int]:
        return list(self.jobs[job_id].ops[op_idx].machine_time.keys())

    def ptime(self, job_id: int, op_idx: int, machine: int) -> int:
        return self.jobs[job_id].ops[op_idx].machine_time[machine]


@dataclass
class ScheduleResult:
    makespan: int
    # dicts: (job_id, op_idx) -> (machine, start, end)
    assignment: Dict[Tuple[int, int], Tuple[int, int, int]]


# ------------------------------
# Example instance (editable)
# ------------------------------
def build_example_problem() -> Problem:
    # Example with 3 jobs and 3 machines. Each operation can run on a subset of machines with different times.
    # You can modify or replace by loading from JSON if you like.
    #
    # Job 0: 3 ops
    j0 = Job(
        ops=[
            Operation({0: 3, 1: 2}),          # Op0: M0=3, M1=2
            Operation({1: 2, 2: 4}),          # Op1: M1=2, M2=4
            Operation({0: 2, 2: 3}),          # Op2: M0=2, M2=3
        ]
    )
    # Job 1: 2 ops
    j1 = Job(
        ops=[
            Operation({0: 2, 2: 1}),          # Op0
            Operation({1: 4, 2: 3}),          # Op1
        ]
    )
    # Job 2: 3 ops
    j2 = Job(
        ops=[
            Operation({0: 2, 1: 3, 2: 2}),    # Op0
            Operation({1: 3}),                # Op1 (only M1 allowed)
            Operation({0: 4, 2: 2}),          # Op2
        ]
    )
    return Problem(jobs=[j0, j1, j2], m=3)


# ------------------------------
# Chromosome representation
# ------------------------------
@dataclass
class Chromosome:
    # job sequence (each job appears as many times as its number of ops)
    job_sequence: List[int]
    # machine choices aligned with the *k-th occurrence* of a job in the sequence
    # We store as a dict: (job_id, occurrence_index) -> machine_id
    machine_choices: Dict[Tuple[int, int], int]
    fitness: Optional[int] = None  # makespan

    def copy(self) -> "Chromosome":
        return Chromosome(self.job_sequence.copy(), self.machine_choices.copy(), self.fitness)


def random_chromosome(pb: Problem) -> Chromosome:
    # Build job_sequence: each job repeated by its number of operations
    seq = []
    job_occ_total = {}
    for j, job in enumerate(pb.jobs):
        k = len(job.ops)
        seq.extend([j] * k)
        job_occ_total[j] = k
    random.shuffle(seq)

    # For machine choices, pick a random feasible machine for each occurrence
    machine_choices = {}
    # We need to map each occurrence to the correct operation index for that job.
    # During decoding, the t-th occurrence corresponds to op_idx=t (0-based).
    for j, k in job_occ_total.items():
        for occ in range(k):
            op_idx = occ
            choices = pb.op_choices(j, op_idx)
            machine_choices[(j, occ)] = random.choice(choices)

    return Chromosome(seq, machine_choices)


# ------------------------------
# Decoder: from chromosome to schedule
# ------------------------------
def decode_makespan(pb: Problem, chrom: Chromosome) -> ScheduleResult:
    # Track next operation index needed per job
    next_op = {j: 0 for j in range(pb.n_jobs)}
    # Machine ready times
    machine_ready = {m: 0 for m in range(pb.m)}
    # Job ready times (due to precedence within job)
    job_ready = {j: 0 for j in range(pb.n_jobs)}
    # Assignment results
    assign: Dict[Tuple[int, int], Tuple[int, int, int]] = {}

    # Counters for occurrences of each job seen in the sequence so far
    occ_seen = {j: 0 for j in range(pb.n_jobs)}

    for job_id in chrom.job_sequence:
        op_idx = next_op[job_id]  # which operation is ready for this job
        if op_idx >= len(pb.jobs[job_id].ops):
            # This job has finished all ops; skip (sequence genes beyond required ops are ignored)
            continue

        occ = occ_seen[job_id]
        occ_seen[job_id] += 1

        # pick machine from stored choice; if infeasible (user changed instance), pick best feasible
        chosen_m = chrom.machine_choices.get((job_id, occ))
        if chosen_m is None or chosen_m not in pb.jobs[job_id].ops[op_idx].machine_time:
            # pick the feasible machine that gives the earliest finish (greedy repair)
            best_m, best_end = None, math.inf
            for m in pb.op_choices(job_id, op_idx):
                start = max(machine_ready[m], job_ready[job_id])
                end = start + pb.ptime(job_id, op_idx, m)
                if end < best_end:
                    best_m, best_end = m, end
            chosen_m = best_m

        start = max(machine_ready[chosen_m], job_ready[job_id])
        pt = pb.ptime(job_id, op_idx, chosen_m)
        end = start + pt

        assign[(job_id, op_idx)] = (chosen_m, start, end)
        machine_ready[chosen_m] = end
        job_ready[job_id] = end
        next_op[job_id] += 1

    makespan = max((t for (_, _, t) in assign.values()), default=0)
    return ScheduleResult(makespan=makespan, assignment=assign)


# ------------------------------
# Genetic Algorithm
# ------------------------------
def tournament_select(pop: List[Chromosome], k: int = 3) -> Chromosome:
    cand = random.sample(pop, k)
    return min(cand, key=lambda c: c.fitness)


def order_crossover(parent1: Chromosome, parent2: Chromosome) -> Tuple[Chromosome, Chromosome]:
    # Order crossover on job sequences (preserve multiplicities)
    n = len(parent1.job_sequence)
    a, b = sorted(random.sample(range(n), 2))
    child1_seq = [None] * n
    child2_seq = [None] * n

    child1_seq[a:b] = parent1.job_sequence[a:b]
    child2_seq[a:b] = parent2.job_sequence[a:b]

    def fill(child_seq, donor_seq, segment):
        missing = []
        # Count needed multiplicity for jobs
        need = {}
        for j in donor_seq:
            need[j] = need.get(j, 0) + 1
        # Reduce counts for already placed
        for j in segment:
            if j is None:
                continue
            need[j] -= 1

        # Build the leftover list in donor order
        for j in donor_seq:
            if need[j] > 0:
                missing.append(j)
                need[j] -= 1

        # Fill None slots in child_seq with missing items in order
        it = iter(missing)
        for i in range(len(child_seq)):
            if child_seq[i] is None:
                child_seq[i] = next(it)

    fill(child1_seq, parent2.job_sequence, child1_seq[a:b])
    fill(child2_seq, parent1.job_sequence, child2_seq[a:b])

    # Machine choices: uniform mixing by (job_id, occurrence)
    def mix_machines(p1: Chromosome, p2: Chromosome) -> Dict[Tuple[int, int], int]:
        keys = set(p1.machine_choices.keys()) | set(p2.machine_choices.keys())
        out = {}
        for k in keys:
            if random.random() < 0.5 and k in p1.machine_choices:
                out[k] = p1.machine_choices[k]
            elif k in p2.machine_choices:
                out[k] = p2.machine_choices[k]
            else:
                # shouldn't happen often; skip
                pass
        return out

    c1 = Chromosome(child1_seq, mix_machines(parent1, parent2))
    c2 = Chromosome(child2_seq, mix_machines(parent1, parent2))
    return c1, c2


def mutate(chrom: Chromosome, pb: Problem, p_swap: float = 0.2, p_machine: float = 0.3) -> None:
    # Swap two positions in job sequence
    if random.random() < p_swap and len(chrom.job_sequence) >= 2:
        i, j = random.sample(range(len(chrom.job_sequence)), 2)
        chrom.job_sequence[i], chrom.job_sequence[j] = chrom.job_sequence[j], chrom.job_sequence[i]

    # Mutate machine choices for a random job occurrence
    if random.random() < p_machine:
        # pick a random job
        j = random.randrange(pb.n_jobs)
        k = sum(1 for _ in pb.jobs[j].ops)
        if k > 0:
            occ = random.randrange(k)
            op_idx = occ
            choices = pb.op_choices(j, op_idx)
            chrom.machine_choices[(j, occ)] = random.choice(choices)


def evaluate(pop: List[Chromosome], pb: Problem) -> None:
    for c in pop:
        res = decode_makespan(pb, c)
        c.fitness = res.makespan


def ga_solve(
    pb: Problem,
    pop_size: int = 80,
    gens: int = 200,
    cx_rate: float = 0.9,
    mut_rate: float = 0.3,
    tourn_k: int = 3,
    verbose: bool = True,
) -> Tuple[Chromosome, ScheduleResult]:
    # init
    pop = [random_chromosome(pb) for _ in range(pop_size)]
    evaluate(pop, pb)
    best = min(pop, key=lambda c: c.fitness).copy()

    if verbose:
        print(f"Init best makespan: {best.fitness}")

    for g in range(gens):
        new_pop: List[Chromosome] = []
        while len(new_pop) < pop_size:
            p1, p2 = tournament_select(pop, tourn_k), tournament_select(pop, tourn_k)
            if random.random() < cx_rate:
                c1, c2 = order_crossover(p1, p2)
            else:
                c1, c2 = p1.copy(), p2.copy()

            if random.random() < mut_rate:
                mutate(c1, pb)
            if random.random() < mut_rate:
                mutate(c2, pb)

            new_pop.extend([c1, c2])

        pop = new_pop[:pop_size]
        evaluate(pop, pb)
        cur_best = min(pop, key=lambda c: c.fitness)
        if cur_best.fitness < best.fitness:
            best = cur_best.copy()

        if verbose and (g + 1) % max(1, gens // 10) == 0:
            print(f"Gen {g+1}/{gens}: best makespan = {best.fitness}")

    best_res = decode_makespan(pb, best)
    return best, best_res


def draw_gantt(pb: Problem, res: ScheduleResult, title: str = "FJSP Schedule") -> None:
    # Prepare bars per machine
    machine_bars = {m: [] for m in range(pb.m)}
    for (j, o), (m, s, e) in res.assignment.items():
        machine_bars[m].append((s, e - s, f"J{j}-O{o}"))

    fig, ax = plt.subplots(figsize=(10, 4))
    y_ticks = []
    y_labels = []
    for idx, m in enumerate(sorted(machine_bars.keys())):
        bars = sorted(machine_bars[m], key=lambda x: x[0])
        for (start, dur, label) in bars:
            ax.barh(idx, dur, left=start)
            ax.text(start + dur / 2, idx, label, va="center", ha="center")
        y_ticks.append(idx)
        y_labels.append(f"M{m}")

    ax.set_yticks(y_ticks)
    ax.set_yticklabels(y_labels)
    ax.set_xlabel("Time")
    ax.set_title(f"{title} (makespan={res.makespan})")
    plt.tight_layout()
    plt.show()


# ------------------------------
# Main (demo)
# ------------------------------
def main():
    pb = build_example_problem()

    # You can also define a JSON file describing the instance.
    # Example JSON structure (jobs list; each job is list of operations; each operation is dict of machine:ptime):
    # [
    #   [ {"0":3, "1":2}, {"1":2, "2":4}, {"0":2, "2":3} ],
    #   [ {"0":2, "2":1}, {"1":4, "2":3} ],
    #   [ {"0":2, "1":3, "2":2}, {"1":3}, {"0":4, "2":2} ]
    # ]
    # To load from a file, replace build_example_problem() with:
    # with open("instance.json","r") as f:
    #     data = json.load(f)
    # jobs = []
    # machines = set()
    # for job_ops in data:
    #     ops = []
    #     for op in job_ops:
    #         op_mt = {int(k): int(v) for k,v in op.items()}
    #         ops.append(Operation(op_mt))
    #         machines.update(op_mt.keys())
    #     jobs.append(Job(ops))
    # pb = Problem(jobs, m=max(machines)+1)

    best, res = ga_solve(pb, pop_size=80, gens=150, verbose=True)
    print(f"Best makespan: {res.makespan}")
    # Pretty print assignment
    for (j, o), (m, s, e) in sorted(res.assignment.items()):
        print(f"Job {j} Op {o} -> M{m}, start={s}, end={e}")

    draw_gantt(pb, res, title="FJSP GA (demo)")


if __name__ == "__main__":
    main()
